import requests
import csv
import time
import random
import pandas as pd
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
import os

# 配置信息
CONFIG = {
    'base_url': 'https://data.stats.gov.cn/search.htm',
    'headers': {
        'Cookie': 'wzws_sessionid=oGhxpkSBZDJhZDY3gDExMi41NC44OS4xNTSCZmM1ZWUx; JSESSIONID=qYL78KFf74e02ojqpehUX7WK7mhz6EzX6VjWbYydvTClhOHLCv0t!-1298377358; u=6; experience=show',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36 Edg/138.0.0.0',
        'x-requested-with': 'XMLHttpRequest'
    },
    'municipalities': ['北京市', '天津市', '上海市', '重庆市'],  # 四个直辖市
    'target_datasource': '分省季度数据',
    'income_keywords': [
        "居民人均可支配收入", "全体居民人均可支配收入",
        "城镇居民人均可支配收入", "农村居民人均可支配收入"
    ],
    'max_workers': 2,  # 并发线程数
    'delay_range': (1, 3)  # 请求延迟范围(秒)
}


# 解析季度数据
def parse_quarter_data(sj):
    """从日期字符串解析年份和季度"""
    if not sj or len(sj) < 5:
        return '', ''
    year = sj[:4]
    quarter_map = {'A': 'Q1', 'B': 'Q2', 'C': 'Q3', 'D': 'Q4'}
    quarter_code = sj[-1]
    return year, quarter_map.get(quarter_code, '')


# 获取指定地区的所有页面数据
def get_all_data(region):
    all_results = []
    page = 0
    search_keyword = f"历年{region}人均收入"

    print(f"\n开始爬取{region}数据，搜索关键词: {search_keyword}")

    while True:
        # 构建请求参数
        params = {
            's': search_keyword,
            'm': 'searchdata',
            'db': '',
            'p': page
        }

        try:
            # 发送请求
            response = requests.get(
                CONFIG['base_url'],
                headers=CONFIG['headers'],
                params=params,
                timeout=10
            )
            response.raise_for_status()  # 检查请求是否成功

            # 解析JSON数据
            data = response.json()

            # 如果没有结果或结果为空，退出循环
            if not data.get('result'):
                print(f"已完成{region}数据爬取，共获取{page}页")
                break

            all_results.extend(data['result'])
            print(f"已获取{region}第 {page + 1} 页数据，共 {len(data['result'])} 条记录")

            # 检查是否还有下一页
            if (page + 1) * 10 >= data.get('pagecount', 0):
                print(f"已完成{region}数据爬取，共获取{page + 1}页")
                break

            page += 1

            # 添加随机延迟，避免频繁请求
            delay = random.uniform(*CONFIG['delay_range'])
            time.sleep(delay)

        except Exception as e:
            print(f"获取{region}第 {page + 1} 页数据失败: {e}")
            # 失败时也增加延迟
            time.sleep(random.uniform(3, 5))
            break

    return all_results, region


# 处理爬取的数据并保存为CSV
def process_crawled_data(data, region):
    if not data:
        print(f"{region}没有数据可处理")
        return None

    # 创建地区目录
    region_dir = f"data/{region}"
    os.makedirs(region_dir, exist_ok=True)

    # 定义CSV文件列名
    fieldnames = ['年份', '季度', '地区', '指标', '数据值', '数据源', '数据链接']

    # 生成CSV文件名
    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
    filename = f"{region_dir}/{region}人均收入数据_{timestamp}.csv"

    with open(filename, 'w', newline='', encoding='utf-8-sig') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

        # 处理每条数据记录
        for record in data:
            # 解析年份和季度
            year, quarter = '', ''
            if 'sj' in record:
                year, quarter = parse_quarter_data(record['sj'])

            # 提取指标名称（去除后缀）
            indicator = record.get('zb', '').split('_')[0]

            # 构建CSV行数据
            row = {
                '年份': year,
                '季度': quarter,
                '地区': record.get('reg', ''),
                '指标': indicator,
                '数据值': record.get('data', ''),
                '数据源': record.get('db', ''),
                '数据链接': f"https://data.stats.gov.cn/easyquery.htm?{record.get('report', '')}"
                if 'report' in record else ''
            }

            # 写入CSV
            writer.writerow(row)

    print(f"{region}原始数据已成功保存到 {filename}")
    print(f"{region}共处理 {len(data)} 条记录")
    return filename


# 按数据源筛选数据
def filter_by_datasource(input_csv, datasource, region):
    """
    按数据源筛选数据

    参数:
        input_csv: 输入CSV文件路径
        datasource: 目标数据源名称
        region: 地区名称
    """
    # 读取CSV数据
    try:
        df = pd.read_csv(input_csv)
    except FileNotFoundError:
        print(f"错误：找不到{region}的输入文件 {input_csv}")
        return None

    # 创建地区目录
    region_dir = f"data/{region}"
    os.makedirs(region_dir, exist_ok=True)

    # 自动生成输出文件名
    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
    output_csv = f"{region_dir}/{datasource.replace('/', '_')}_数据_{timestamp}.csv"

    # 筛选出特定数据源的记录
    filtered_df = df[df['数据源'] == datasource].copy()

    # 保存筛选后的数据
    if not filtered_df.empty:
        filtered_df.to_csv(output_csv, index=False, encoding='utf-8-sig')
        print(f"{region}已筛选出 {len(filtered_df)} 条来自 '{datasource}' 的数据")
        print(f"{region}筛选后的数据已保存到: {output_csv}")
        return output_csv
    else:
        print(f"{region}未找到来自 '{datasource}' 的数据")
        return None


# 筛选居民人均可支配收入数据
def filter_income_data(input_csv, region):
    """
    筛选居民人均可支配收入数据并去除数据链接字段

    参数:
        input_csv: 输入CSV文件路径
        region: 地区名称
    """
    # 读取CSV数据
    try:
        df = pd.read_csv(input_csv)
    except FileNotFoundError:
        print(f"错误：找不到{region}的输入文件 {input_csv}")
        return None

    # 创建地区目录
    region_dir = f"data/{region}"
    os.makedirs(region_dir, exist_ok=True)

    # 自动生成输出文件名
    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
    output_csv = f"{region_dir}/{region}居民人均可支配收入数据_{timestamp}.csv"

    # 筛选出指标包含"居民人均可支配收入"的记录
    mask = df['指标'].str.contains('|'.join(CONFIG['income_keywords']), na=False)
    filtered_df = df[mask].copy()

    # 从数据链接提取季度信息（如果季度为空）
    def extract_quarter_from_link(link):
        if pd.notna(link) and 'sj=' in link:
            sj_part = link.split('sj=')[1].split('&')[0]
            if len(sj_part) >= 5:
                quarter_code = sj_part[-1]
                quarter_map = {'A': '第一季度', 'B': '第二季度', 'C': '第三季度', 'D': '第四季度'}
                return quarter_map.get(quarter_code, '')
        return ''

    # 填充空的季度信息
    if '季度' in filtered_df.columns and '数据链接' in filtered_df.columns:
        filtered_df.loc[filtered_df['季度'].isna() | (filtered_df['季度'] == ''), '季度'] = \
            filtered_df['数据链接'].apply(extract_quarter_from_link)

    # 移除数据链接列
    if '数据链接' in filtered_df.columns:
        filtered_df = filtered_df.drop(columns=['数据链接'])

    # 调整列顺序
    desired_columns = ['年份', '季度', '地区', '指标', '数据值', '数据源']
    filtered_df = filtered_df.reindex(columns=desired_columns)

    # 保存处理后的数据
    filtered_df.to_csv(output_csv, index=False, encoding='utf-8-sig')
    print(f"{region}已筛选出 {len(filtered_df)} 条居民人均可支配收入数据")
    print(f"{region}处理后的数据已保存到: {output_csv}")
    return output_csv


# 处理单个地区的数据
def process_region(region):
    try:
        # 爬取数据
        data, region = get_all_data(region)

        if data:
            # 处理爬取的数据并保存为CSV
            raw_data_file = process_crawled_data(data, region)

            if raw_data_file:
                # 按数据源筛选
                datasource_filtered_file = filter_by_datasource(
                    raw_data_file,
                    CONFIG['target_datasource'],
                    region
                )

                if datasource_filtered_file:
                    # 筛选居民人均可支配收入数据
                    filter_income_data(datasource_filtered_file, region)

        return f"{region}数据处理完成"

    except Exception as e:
        print(f"处理{region}数据时发生错误: {e}")
        return f"{region}数据处理失败"


def main():
    print("开始爬取四个直辖市的国家统计局数据...")

    # 创建数据目录
    os.makedirs('data', exist_ok=True)

    # 使用线程池并发处理多个地区
    with ThreadPoolExecutor(max_workers=CONFIG['max_workers']) as executor:
        results = list(executor.map(process_region, CONFIG['municipalities']))

    # 打印结果摘要
    print("\n===== 处理结果摘要 =====")
    for result in results:
        print(result)

    print("所有数据爬取与处理完成！")


if __name__ == "__main__":
    main()